import json
import os
from tqdm import tqdm
import pandas as pd

path_dir = os.path.dirname(os.path.abspath(__file__))

if __name__ == '__main__':
    phase25_model = 'qwen3-max'
    with open(os.path.join(path_dir, "phase2.5_output", f"{phase25_model}_prop_desc.jsonl"), 'r',
              encoding='utf-8') as fp:
        lines = fp.readlines()
        prop_descs_data = [json.loads(line) for line in lines]

    phase3_model = 'qwen3-max'
    with open(os.path.join(path_dir, "phase3_output", f"{phase3_model}_questions.jsonl"), 'r', encoding='utf-8') as fp:
        lines = fp.readlines()
        questions_data = [json.loads(line) for line in lines]

    # 1. 所有的desc，构建成文档，一个是多份的，一个是单份的
    all_desc = []
    for idx, prop_desc_data in tqdm(enumerate(prop_descs_data)):
        with open(os.path.join(path_dir, "phase4_output", "docs", f"{idx}.txt"), 'w', encoding='utf-8') as fp:
            all_desc.append(prop_desc_data['desc'])
            fp.write(prop_desc_data['desc'])
    with open(os.path.join(path_dir, "phase4_output", f"{phase3_model}_descriptions.json"), 'w',
              encoding='utf-8') as fp:
        json.dump(all_desc, fp, ensure_ascii=False, indent=2)

    # 2. 所有的ref，构建成文档
    all_ref = []
    cnt = 0
    for idx, prop_desc_data in tqdm(enumerate(prop_descs_data)):
        if prop_desc_data['define_with_ref']:
            with open(os.path.join(path_dir, "phase4_output", "refs", f"{idx}.txt"), 'w', encoding='utf-8') as fp:
                all_ref.append(prop_desc_data['ref']['content'])
                fp.write(prop_desc_data['ref']['content'])
    with open(os.path.join(path_dir, "phase4_output", f"{phase3_model}_references.json"), 'w', encoding='utf-8') as fp:
        json.dump(all_ref, fp, ensure_ascii=False, indent=2)

    # 3. 所有的question，构建成文档
    output_fp = open(os.path.join(path_dir, "phase4_output", f"{phase3_model}_questions.jsonl"), 'w', encoding='utf-8')
    for question_data in tqdm(questions_data):
        ques_str = question_data['answer']
        ques_str = ques_str.strip('`')  # 去掉json
        ques_str = ques_str[4:]  # 去掉json
        questions = json.loads(ques_str)
        for q in questions:
            q['meta'] = question_data['prop']
            if q['rule'] == 'RULE-A':
                q['type'] = '定义题'
                # q['demand_point'] = '规定时间内完成,回答不包含过时信息,应答尽答,禁止答非所问,回答要包含引用,回答包含全部信息,标准答案完全一致'
            elif q['rule'] == 'RULE-B':
                q['type'] = '自由发挥题'
                # q['demand_point'] = '规定时间内完成,回答不包含事实性知识错误,回答不包含逻辑错误,回答不包含过时信息,应答尽答,禁止答非所问,回答要包含引用,回答包含全部信息'
            elif q['rule'] == 'RULE-C':
                q['type'] = '自定义题'
                # q['demand_point'] = ''
            else:
                raise ValueError('rule should be RULE-A, RULE-B, RULE-C')
            output_fp.write(json.dumps(q, ensure_ascii=False) + '\n')
    output_fp.close()

    df = pd.read_json(os.path.join(path_dir, "phase4_output", f"{phase3_model}_questions.jsonl"), lines=True)
    df.to_excel(os.path.join(path_dir, "phase4_output", f"{phase3_model}_questions.xlsx"))
